{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from Models import Transformer, Encoder, Decoder\n",
    "\n",
    "class myModel(nn.Module):\n",
    "    def __init__(self, vocab_size, src, d_model, N, heads, dropout, n_classes):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(vocab_size, d_model, N, heads, dropout)\n",
    "        self.linear1 = nn.Linear(d_model, 1)\n",
    "        self.linear2 = nn.Linear(src, n_classes)\n",
    "        self.fc = nn.Sigmoid()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x, None)\n",
    "        x = self.linear1(x)\n",
    "        x = x.squeeze(-1)\n",
    "        x = self.linear2(x)\n",
    "        x = self.fc(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "# from sklearn import datasets\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import TensorDataset\n",
    "\n",
    "\n",
    "class KddData(object):\n",
    "\n",
    "    def __init__(self, batch_size, file_name, augmented=False, test_filename=None):\n",
    "        self._encoder = {\n",
    "            'label':    LabelEncoder()\n",
    "        }\n",
    "        self.batch_size = batch_size\n",
    "        if (augmented == False):\n",
    "            pass\n",
    "        else:\n",
    "            data = pd.read_csv(file_name)\n",
    "            target = np.array(data['label'])\n",
    "            features = np.array(data.drop('label', axis=1))\n",
    "            X_train, y_train = self.__encode_data(features, target)\n",
    "            self.train_dataset = TensorDataset(\n",
    "                torch.from_numpy(X_train.astype(np.float32)),\n",
    "                torch.from_numpy(y_train.astype(np.int))\n",
    "            )\n",
    "            test_data = pd.read_csv(test_filename)\n",
    "            target = np.array(test_data['label'])\n",
    "            print(target.shape)\n",
    "            features = np.array(test_data.drop('label', axis=1))\n",
    "            X_test, y_test = self.__encode_data(features, target)\n",
    "            self.test_dataset = TensorDataset(\n",
    "                torch.from_numpy(X_test.astype(np.float32)),\n",
    "                torch.from_numpy(y_test.astype(np.float32))\n",
    "            )\n",
    "            self.train_dataloader = DataLoader(self.train_dataset, self.batch_size, shuffle=True)\n",
    "            self.test_dataloader = DataLoader(self.test_dataset, self.batch_size, shuffle=True)\n",
    "            return\n",
    "        data = pd.read_csv(file_name)\n",
    "        target = np.array(data['label'])\n",
    "        features = np.array(data.drop('label', axis=1))\n",
    "\n",
    "        data_X, data_y = self.__encode_data(features, target)\n",
    "        self.train_dataset, self.test_dataset = self.__split_data_to_tensor(data_X, data_y)\n",
    "        self.train_dataloader = DataLoader(self.train_dataset, self.batch_size, shuffle=True)\n",
    "        self.test_dataloader = DataLoader(self.test_dataset, self.batch_size, shuffle=True)\n",
    "\n",
    "    \"\"\"将数据中字符串部分转换为数字，并将输入的41维特征转换为8*8的矩阵\"\"\"\n",
    "    def __encode_data(self, data_X, data_y):\n",
    "        self._encoder['label'].fit(list(set(data_y)))\n",
    "        data_X = np.pad(data_X, ((0, 0), (0, 100 - len(data_X[0]))), 'constant').reshape(-1, 1, 10, 10)\n",
    "        data_y = self._encoder['label'].transform(data_y)\n",
    "        return data_X, data_y\n",
    "\n",
    "    \"\"\"将数据拆分为训练集和测试集，并转换为TensorDataset对象\"\"\"\n",
    "    def __split_data_to_tensor(self, data_X, data_y):\n",
    "        X_train, X_test, y_train, y_test = train_test_split(data_X, data_y, test_size=0.3)\n",
    "        train_dataset = TensorDataset(\n",
    "            torch.from_numpy(X_train.astype(np.float32)),\n",
    "            torch.from_numpy(y_train.astype(np.int))\n",
    "        )\n",
    "        test_dataset = TensorDataset(\n",
    "            torch.from_numpy(X_test.astype(np.float32)),\n",
    "            torch.from_numpy(y_test.astype(np.int))\n",
    "        )\n",
    "        return train_dataset, test_dataset\n",
    "\n",
    "    \"\"\"接受一个数组进行解码\"\"\"\n",
    "    def decode(self, data, label=False):\n",
    "        if not label:\n",
    "            _data = list(data)\n",
    "            _data[1] = self._encoder['protocal'].inverse_transform([_data[1]])[0]\n",
    "            _data[2] = self._encoder['service'].inverse_transform([_data[2]])[0]\n",
    "            _data[2] = self._encoder['flag'].inverse_transform([_data[3]])[0]\n",
    "            return _data\n",
    "        return self._encoder['label'].inverse_transform(data)\n",
    "    \n",
    "    def encode(self, data, label=False):\n",
    "        if not label:\n",
    "            _data = list(data)\n",
    "            _data[1] = self._encoder['protocal'].transform([_data[1]])[0]\n",
    "            _data[2] = self._encoder['service'].transform([_data[2]])[0]\n",
    "            _data[3] = self._encoder['flag'].transform([_data[3]])[0]\n",
    "            return _data\n",
    "        return self._encoder['label'].transform([data])[0]\n",
    "\n",
    "batch_size = 128\n",
    "\n",
    "# dataset.train_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "# 神经网络参数\n",
    "batch_size = 128\n",
    "learning_rate = 5e-3\n",
    "num_epoches = 30\n",
    "USE_GPU = torch.cuda.is_available()\n",
    "\n",
    "data_file = \"E:\\THU\\科研\\SRT\\cross_network_defense\\model_reproduce\\dataset\\\\train-expon-nagle-open.csv\"\n",
    "test_file = \"E:\\THU\\科研\\SRT\\cross_network_defense\\model_reproduce\\dataset\\\\test-lossRate-0.csv\"\n",
    "dataset = KddData(batch_size, file_name=data_file, augmented=True, test_filename=test_file)\n",
    "model = myModel(1500, 100, 8, 6, 8, 0.1, 1)\n",
    "model = model.cuda()\n",
    "criterion = nn.BCELoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for epoch in range(num_epoches):\n",
    "    print('epoch {}'.format(epoch + 1))\n",
    "    print('*' * 10)\n",
    "    running_loss = 0.0\n",
    "    running_acc = 0.0\n",
    "    for i, data in enumerate(dataset.train_dataloader, 1):\n",
    "        model.train()\n",
    "        x, label = data\n",
    "        # Reshape x and convert to LongTensor\n",
    "        x = x.reshape(x.size(0),  100)\n",
    "        x = Variable(x)\n",
    "        x = x.long()\n",
    "        #print(x.max())\n",
    "        #print(x.min())\n",
    "        x = x.cuda()\n",
    "        label = label.reshape(label.size(0), 1)\n",
    "        label = label.float()\n",
    "        label = label.cuda()\n",
    "        out = model(x)\n",
    "        #print(out)\n",
    "        #print(label)\n",
    "        loss = criterion(out, label)\n",
    "        running_loss += loss.item() * label.size(0)\n",
    "        pred = torch.round(out)\n",
    "        num_correct = (pred == label).sum()\n",
    "        #print(\"num_correct = \" + str(num_correct))\n",
    "        accuracy = (pred == label).float().mean()\n",
    "        running_acc += num_correct.item()\n",
    "            \n",
    "        # Back Propagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5 (default, Sep  3 2020, 21:29:08) [MSC v.1916 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e42634819b8c191a5d07eaf23810ff32516dd8d3875f28ec3e488928fbd3c187"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
